"""
黑方控制模块
处理黑方无人艇的AI行为和决策逻辑

作者：wjttdbx
版本：2.0
"""
import numpy as np
import math
from typing import List, Tuple, Optional
import sys
sys.path.insert(0, sys.path[0]+"/../")
from core.config import GameConfig
from core.units import BlackBoat, WhiteBoat, WhiteDrone
from sh15.core.utils import point_in_polygon
from sh15.core.BoundarySystem import BoundarySystem


class BlackAgent:
    """
    黑方AI智能体基类
    
    定义了黑方AI需要实现的核心接口
    """
    
    def update_all_black_boats(self, black_boats: List[BlackBoat], white_boats: List[WhiteBoat], 
                               current_time: float, dt: float):
        """
        更新所有黑方无人艇的状态
        
        Args:
            black_boats (List[BlackBoat]): 黑方无人艇列表
            white_boats (List[WhiteBoat]): 白方无人艇列表
            current_time (float): 当前时间
            dt (float): 时间步长
        """
        raise NotImplementedError
    
    def should_attempt_lock(self, black_boat: BlackBoat, white_boats: List[WhiteBoat]) -> bool:
        """
        判断是否应该尝试锁定
        
        Args:
            black_boat (BlackBoat): 当前黑方无人艇
            white_boats (List[WhiteBoat]): 白方无人艇列表
            
        Returns:
            bool: 是否尝试锁定
        """
        return False
        
    def reset(self):
        """
        重置智能体状态，用于新回合开始
        """
        pass


class RuleBasedBlackAgent(BlackAgent):
    """
    基于规则的高级黑方AI智能体
    
    包含更复杂的策略和战术：
    1. 动态编队变换
    2. 智能分工合作
    3. 多阶段突破策略
    4. 实时战术调整
    """
    
    def __init__(self):
        """初始化高级黑方AI"""
        super().__init__()
        self.formation_type = "adaptive"  # 编队类型：adaptive, spread, wedge, line
        self.coordination_enabled = True  # 是否启用协调
        self.breakthrough_phase = "approach"  # 突破阶段：approach, penetrate, rush
        self.team_strategy = "balanced"  # 团队策略：balanced, sacrifice, conservative
    
    def update_all_black_boats(self, black_boats: List[BlackBoat], white_boats: List[WhiteBoat], 
                              current_time: float, dt: float):
        """
        高级AI的整体更新逻辑
        
        重写基类方法，加入全局战术协调
        """
        # 1. 全局态势评估
        self._update_global_strategy(black_boats, white_boats, current_time)
        
        # 2. 为每艘船分配角色
        ship_roles = self._assign_ship_roles(black_boats, white_boats)
        
        # 3. 根据角色更新每艘船
        for i, black_boat in enumerate(black_boats):
            if not black_boat.active or black_boat.is_frozen(current_time):
                continue
            
            ship_role = ship_roles.get(black_boat.id, "advance")
            self._update_boat_with_role(black_boat, ship_role, black_boats, white_boats, current_time, dt)
    
    def _update_global_strategy(self, black_boats: List[BlackBoat], white_boats: List[WhiteBoat], current_time: float):
        """更新全局策略"""
        active_black = sum(1 for boat in black_boats if boat.active)
        active_white = sum(1 for boat in white_boats if boat.active)
        
        # 计算平均推进距离
        avg_progress = sum(300000 - boat.x for boat in black_boats if boat.active) / max(active_black, 1)
        
        # 根据战况调整策略
        if avg_progress > 250000:  # 接近突破线
            self.breakthrough_phase = "rush"
            self.team_strategy = "sacrifice"  # 牺牲策略，一部分船吸引火力
        elif avg_progress > 150000:  # 中距离
            self.breakthrough_phase = "penetrate"
            self.team_strategy = "balanced"
        else:  # 远距离
            self.breakthrough_phase = "approach"
            self.team_strategy = "conservative"
    
    def _assign_ship_roles(self, black_boats: List[BlackBoat], white_boats: List[WhiteBoat]) -> dict:
        """为每艘船分配战术角色"""
        roles = {}
        active_boats = [boat for boat in black_boats if boat.active]
        
        if not active_boats:
            return roles
        
        # 根据当前策略分配角色
        if self.team_strategy == "sacrifice":
            # 牺牲策略：部分船作为诱饵
            for i, boat in enumerate(active_boats):
                if i < len(active_boats) // 3:
                    roles[boat.id] = "decoy"  # 诱饵
                else:
                    roles[boat.id] = "breakthrough"  # 突破
        
        elif self.team_strategy == "balanced":
            # 平衡策略：协调推进
            for i, boat in enumerate(active_boats):
                if i % 2 == 0:
                    roles[boat.id] = "advance"  # 推进
                else:
                    roles[boat.id] = "support"  # 支援
        
        else:  # conservative
            # 保守策略：谨慎推进
            for boat in active_boats:
                roles[boat.id] = "cautious_advance"
        
        return roles
    
    def _update_boat_with_role(self, black_boat: BlackBoat, role: str, all_black_boats: List[BlackBoat],
                              white_boats: List[WhiteBoat], current_time: float, dt: float):
        """根据角色更新船只行为"""
        if role == "decoy":
            # 诱饵：主动吸引白方注意力
            target_x, target_y = self._execute_decoy_behavior(black_boat, white_boats)
        elif role == "breakthrough":
            # 突破：专注突破，避免战斗
            target_x, target_y = self._execute_stealth_breakthrough(black_boat, white_boats)
        elif role == "support":
            # 支援：协助其他船只
            target_x, target_y = self._execute_support_behavior(black_boat, all_black_boats, white_boats)
        elif role == "cautious_advance":
            # 谨慎推进：小心前进
            target_x, target_y = self._execute_cautious_advance(black_boat, white_boats)
        else:  # advance
            # 正常推进
            target_x, target_y = self._execute_tactical_advance(black_boat, white_boats, all_black_boats)
        
        # 应用群体协调
        target_x, target_y = self._apply_formation_keeping(black_boat, all_black_boats, target_x, target_y)
        
        # 执行移动
        self._move_towards_target(black_boat, target_x, target_y, dt)
        
        # 边界约束和突破检查
        BoundarySystem.enforce_boundaries(black_boat)
        BoundarySystem.check_breakthrough(black_boat)
    
    def _execute_decoy_behavior(self, black_boat: BlackBoat, white_boats: List[WhiteBoat]) -> Tuple[float, float]:
        """执行诱饵行为"""
        # 主动接近白方，吸引注意力
        closest_white = min(white_boats, key=lambda wb: black_boat.distance_to(wb) if wb.active else float('inf'))
        if closest_white and closest_white.active:
            # 在锁定范围边缘移动，吸引但不被轻易锁定
            distance = black_boat.distance_to(closest_white)
            if distance > GameConfig.WHITE_BOAT_LOCK_RANGE:
                # 接近白方
                return closest_white.x, closest_white.y
            else:
                # 在锁定范围边缘绕行
                angle = math.atan2(black_boat.y - closest_white.y, black_boat.x - closest_white.x)
                angle += math.pi / 4  # 转向45度
                target_distance = GameConfig.WHITE_BOAT_LOCK_RANGE * 1.1
                return (closest_white.x + math.cos(angle) * target_distance,
                       closest_white.y + math.sin(angle) * target_distance)
        
        return self._execute_rush_breakthrough(black_boat)
    
    def _execute_stealth_breakthrough(self, black_boat: BlackBoat, white_boats: List[WhiteBoat]) -> Tuple[float, float]:
        """执行隐蔽突破"""
        # 寻找最隐蔽的路径
        gap_x, gap_y = self._find_defensive_gap(black_boat, white_boats)
        if gap_x is not None:
            return gap_x, gap_y
        
        # 如果没有缺口，选择最远离白方的路径
        best_path = self._find_least_defended_path(black_boat, white_boats)
        return best_path
    
    def _execute_support_behavior(self, black_boat: BlackBoat, all_black_boats: List[BlackBoat],
                                 white_boats: List[WhiteBoat]) -> Tuple[float, float]:
        """执行支援行为"""
        # 寻找需要支援的友军
        threatened_allies = []
        for ally in all_black_boats:
            if ally.id == black_boat.id or not ally.active:
                continue
            
            threat_level = self._assess_local_threat(ally, white_boats)
            if threat_level > 0.6:
                threatened_allies.append((ally, threat_level))
        
        if threatened_allies:
            # 支援威胁最大的友军
            ally_to_help = max(threatened_allies, key=lambda x: x[1])[0]
            # 移动到友军和敌人之间
            closest_enemy = min(white_boats, key=lambda wb: ally_to_help.distance_to(wb) if wb.active else float('inf'))
            if closest_enemy and closest_enemy.active:
                support_x = (ally_to_help.x + closest_enemy.x) / 2
                support_y = (ally_to_help.y + closest_enemy.y) / 2
                return support_x, support_y
        
        # 没有需要支援的友军，正常推进
        return self._execute_tactical_advance(black_boat, white_boats, all_black_boats)
    
    def _execute_cautious_advance(self, black_boat: BlackBoat, white_boats: List[WhiteBoat]) -> Tuple[float, float]:
        """执行谨慎推进"""
        # 以较慢的速度推进，优先避险
        target_x, target_y = self._execute_rush_breakthrough(black_boat)
        
        # 强化避障
        avoidance_x, avoidance_y = self._apply_avoidance(black_boat, white_boats, target_x, target_y)
        
        # 保守推进：只前进一半距离
        progress_x = black_boat.x + (avoidance_x - black_boat.x) * 0.5
        progress_y = black_boat.y + (avoidance_y - black_boat.y) * 0.5
        
        return progress_x, progress_y
    
    def _find_least_defended_path(self, black_boat: BlackBoat, white_boats: List[WhiteBoat]) -> Tuple[float, float]:
        """寻找防守最薄弱的路径"""
        # 评估不同Y坐标的防守强度
        line_y_start = 0
        line_y_end = 600000
        
        best_y = black_boat.y
        min_defense = float('inf')
        
        # 测试不同的Y坐标
        for test_y in range(int(line_y_start), int(line_y_end), 10000):
            defense_strength = 0
            for white_boat in white_boats:
                if not white_boat.active:
                    continue
                
                # 计算到测试路径的距离
                path_distance = abs(white_boat.y - test_y)
                if path_distance < GameConfig.WHITE_BOAT_DETECT_RANGE:
                    defense_strength += (GameConfig.WHITE_BOAT_DETECT_RANGE - path_distance) / GameConfig.WHITE_BOAT_DETECT_RANGE
            
            if defense_strength < min_defense:
                min_defense = defense_strength
                best_y = test_y
        
        return 0, best_y  # 朝突破线的最佳Y坐标移动

    def _execute_rush_breakthrough(self, black_boat: BlackBoat) -> Tuple[float, float]:
        """执行冲刺突破"""
        # 直接朝最近的突破点冲刺
        target_x = 0  # A1A6线
        target_y = black_boat.y  # 保持当前Y坐标
        return target_x, target_y
    
    def _assess_threat_level(self, black_boat: BlackBoat, white_boats: List[WhiteBoat]) -> float:
        """
        评估当前威胁等级
        
        Returns:
            float: 威胁等级 (0-1)
        """
        threat_score = 0.0
        
        for white_boat in white_boats:
            if not white_boat.active:
                continue
                
            distance = black_boat.distance_to(white_boat)
            max_threat_range = GameConfig.WHITE_BOAT_LOCK_RANGE
            
            if distance <= max_threat_range:
                # 距离越近威胁越大
                distance_threat = (max_threat_range - distance) / max_threat_range
                threat_score += distance_threat
        
        return min(threat_score, 1.0)
    
    def _execute_dispersion_strategy(self, black_boat: BlackBoat, 
                                   all_black_boats: List[BlackBoat]) -> Tuple[float, float]:
        """执行分散策略"""
        # 计算群体中心
        active_boats = [boat for boat in all_black_boats if boat.active and boat.id != black_boat.id]
        if not active_boats:
            return 0.0, 0.0
        
        center_x = sum(boat.x for boat in active_boats) / len(active_boats)
        center_y = sum(boat.y for boat in active_boats) / len(active_boats)
        
        # 远离群体中心
        dispersion_x = black_boat.x - center_x
        dispersion_y = black_boat.y - center_y
        
        # 标准化并放大
        distance = math.sqrt(dispersion_x**2 + dispersion_y**2)
        if distance > 0:
            dispersion_x = (dispersion_x / distance) * 10000  # 10km分散距离
            dispersion_y = (dispersion_y / distance) * 10000
        
        return dispersion_x, dispersion_y
    
    def _execute_evasion_strategy(self, black_boat: BlackBoat, 
                                white_boats: List[WhiteBoat]) -> Tuple[float, float]:
        """执行规避策略"""
        evasion_x = 0.0
        evasion_y = 0.0
        
        for white_boat in white_boats:
            if not white_boat.active:
                continue
                
            distance = black_boat.distance_to(white_boat)
            if distance < GameConfig.WHITE_BOAT_LOCK_RANGE * 1.5:
                # 计算规避方向
                dx = black_boat.x - white_boat.x
                dy = black_boat.y - white_boat.y
                distance_norm = max(distance, 1.0)
                
                evasion_strength = (GameConfig.WHITE_BOAT_LOCK_RANGE * 1.5 - distance) / (GameConfig.WHITE_BOAT_LOCK_RANGE * 1.5)
                evasion_x += (dx / distance_norm) * evasion_strength * 15000
                evasion_y += (dy / distance_norm) * evasion_strength * 15000
        
        return evasion_x, evasion_y
    
    def _execute_rush_strategy(self, black_boat: BlackBoat) -> Tuple[float, float]:
        """执行快速突破策略"""
        # 直接朝突破线冲刺
        rush_x = -20000  # 强力朝X=0方向移动
        rush_y = 0.0     # 保持Y方向
        
        return rush_x, rush_y

    def _execute_tactical_advance(self, black_boat: BlackBoat, white_boats: List[WhiteBoat], 
                                all_black_boats: List[BlackBoat]) -> Tuple[float, float]:
        """执行战术推进"""
        # 获取基础突破目标
        target_x, target_y = self._execute_rush_breakthrough(black_boat)
        
        # 应用避障
        target_x, target_y = self._apply_avoidance(black_boat, white_boats, target_x, target_y)
        
        # 寻找薄弱环节
        weak_point_x, weak_point_y = self._find_defensive_gap(black_boat, white_boats)
        if weak_point_x is not None:
            # 朝薄弱环节移动
            target_x = weak_point_x
            target_y = weak_point_y
        
        return target_x, target_y

    def _apply_avoidance(self, black_boat: BlackBoat, white_boats: List[WhiteBoat], 
                        target_x: float, target_y: float) -> Tuple[float, float]:
        """
        应用避障逻辑
        
        Args:
            black_boat (BlackBoat): 黑方无人艇
            white_boats (List[WhiteBoat]): 白方无人艇列表
            target_x (float): 原始目标X坐标
            target_y (float): 原始目标Y坐标
            
        Returns:
            Tuple[float, float]: 调整后的目标位置
        """
        avoidance_force_x = 0.0
        avoidance_force_y = 0.0
        
        for white_boat in white_boats:
            if not white_boat.active:
                continue
                
            distance = black_boat.distance_to(white_boat)
            if distance < GameConfig.BLACK_BOAT_DETECT_RANGE:
                # 计算避障力
                # 距离越近，避障力越强
                avoidance_strength = (GameConfig.BLACK_BOAT_DETECT_RANGE - distance) / GameConfig.BLACK_BOAT_DETECT_RANGE
                avoidance_strength = min(avoidance_strength, 1.0)
                
                # 计算避障方向（远离白方无人艇）
                dx = black_boat.x - white_boat.x
                dy = black_boat.y - white_boat.y
                distance_norm = max(distance, 1.0)  # 避免除零
                
                # 避障力强度
                avoidance_multiplier = 20000  # 20km的避障距离
                avoidance_force_x += (dx / distance_norm) * avoidance_strength * avoidance_multiplier
                avoidance_force_y += (dy / distance_norm) * avoidance_strength * avoidance_multiplier
        
        # 应用避障力到目标位置
        adjusted_target_x = target_x + avoidance_force_x
        adjusted_target_y = target_y + avoidance_force_y
        
        return adjusted_target_x, adjusted_target_y

    def _find_defensive_gap(self, black_boat: BlackBoat, white_boats: List[WhiteBoat]) -> Tuple[Optional[float], Optional[float]]:
        """寻找白方防线的薄弱环节"""
        # 将防线分成几个区域，寻找无人艇密度最低的区域
        line_y_start = 50000 * math.sqrt(3)
        line_y_end = 500000 + 50000 * math.sqrt(3)
        
        # 分成5个区域
        best_gap_y = None
        min_defenders = float('inf')
        
        for i in range(5):
            region_y_start = line_y_start + (line_y_end - line_y_start) * i / 5
            region_y_end = line_y_start + (line_y_end - line_y_start) * (i + 1) / 5
            region_y_center = (region_y_start + region_y_end) / 2
            
            # 计算这个区域的白方无人艇数量
            defenders_in_region = 0
            for white_boat in white_boats:
                if not white_boat.active:
                    continue
                if region_y_start <= white_boat.y <= region_y_end:
                    defenders_in_region += 1
            
            if defenders_in_region < min_defenders:
                min_defenders = defenders_in_region
                best_gap_y = region_y_center
        
        if best_gap_y is not None:
            return 10000, best_gap_y  # 朝薄弱区域前进
        
        return None, None
    
    def _apply_formation_keeping(self, black_boat: BlackBoat, all_black_boats: List[BlackBoat],
                               target_x: float, target_y: float) -> Tuple[float, float]:
        """
        应用群体协调逻辑（保持适当分散，避免聚集）

        Args:
            black_boat (BlackBoat): 当前黑方无人艇
            all_black_boats (List[BlackBoat]): 所有黑方无人艇
            target_x (float): 目标X坐标
            target_y (float): 目标Y坐标

        Returns:
            Tuple[float, float]: 调整后的目标位置
        """
        separation_force_x = 0.0
        separation_force_y = 0.0

        separation_distance = 5000  # 5km分离距离

        for other_boat in all_black_boats:
            if other_boat.id == black_boat.id or not other_boat.active:
                continue

            distance = black_boat.distance_to(other_boat)
            if separation_distance > distance > 0:
                # 计算分离力
                separation_strength = (separation_distance - distance) / separation_distance

                # 计算分离方向
                dx = black_boat.x - other_boat.x
                dy = black_boat.y - other_boat.y
                distance_norm = max(distance, 1.0)

                # 分离力强度
                separation_multiplier = 2000  # 2km的分离力
                separation_force_x += (dx / distance_norm) * separation_strength * separation_multiplier
                separation_force_y += (dy / distance_norm) * separation_strength * separation_multiplier

        # 应用分离力
        adjusted_target_x = target_x + separation_force_x
        adjusted_target_y = target_y + separation_force_y

        return adjusted_target_x, adjusted_target_y
    
    def _move_towards_target(self, black_boat: BlackBoat, target_x: float, target_y: float, dt: float):
        dx = target_x - black_boat.x
        dy = target_y - black_boat.y
        distance = math.sqrt(dx*dx + dy*dy)
        if distance > 0:
            heading = math.atan2(dy, dx)
            speed = black_boat.max_speed
            lock_action = 0.0
            drone_action = 0.0
            black_boat.move_with_kinematics(speed, heading, dt)
        else:
            speed = 0.0
            heading = black_boat.heading
            lock_action = 0.0
            drone_action = 0.0
            black_boat.move_with_kinematics(speed, heading, dt)

    def _assess_local_threat(self, black_boat: BlackBoat, white_boats: List[WhiteBoat]) -> float:
        """评估局部威胁等级"""
        threat_score = 0.0
        
        for white_boat in white_boats:
            if not white_boat.active:
                continue
            
            distance = black_boat.distance_to(white_boat)
            
            # 距离威胁
            if distance <= GameConfig.WHITE_BOAT_LOCK_RANGE:
                threat_score += 1.0  # 锁定范围内，最高威胁
            elif distance <= GameConfig.WHITE_BOAT_DETECT_RANGE:
                threat_score += 0.6  # 探测范围内，中等威胁
            elif distance <= 50000:  # 50km内
                threat_score += 0.3  # 潜在威胁
        
        return min(threat_score, 1.0)
    
    def _execute_dispersion_strategy(self, black_boat: BlackBoat, 
                                   all_black_boats: List[BlackBoat]) -> Tuple[float, float]:
        """执行分散策略"""
        # 计算群体中心
        active_boats = [boat for boat in all_black_boats if boat.active and boat.id != black_boat.id]
        if not active_boats:
            return 0.0, 0.0
        
        center_x = sum(boat.x for boat in active_boats) / len(active_boats)
        center_y = sum(boat.y for boat in active_boats) / len(active_boats)
        
        # 远离群体中心
        dispersion_x = black_boat.x - center_x
        dispersion_y = black_boat.y - center_y
        
        # 标准化并放大
        distance = math.sqrt(dispersion_x**2 + dispersion_y**2)
        if distance > 0:
            dispersion_x = (dispersion_x / distance) * 10000  # 10km分散距离
            dispersion_y = (dispersion_y / distance) * 10000
        
        return dispersion_x, dispersion_y
    
    def _execute_evasion_strategy(self, black_boat: BlackBoat, 
                                white_boats: List[WhiteBoat]) -> Tuple[float, float]:
        """执行规避策略"""
        evasion_x = 0.0
        evasion_y = 0.0
        
        for white_boat in white_boats:
            if not white_boat.active:
                continue
                
            distance = black_boat.distance_to(white_boat)
            if distance < GameConfig.WHITE_BOAT_LOCK_RANGE * 1.5:
                # 计算规避方向
                dx = black_boat.x - white_boat.x
                dy = black_boat.y - white_boat.y
                distance_norm = max(distance, 1.0)
                
                evasion_strength = (GameConfig.WHITE_BOAT_LOCK_RANGE * 1.5 - distance) / (GameConfig.WHITE_BOAT_LOCK_RANGE * 1.5)
                evasion_x += (dx / distance_norm) * evasion_strength * 15000
                evasion_y += (dy / distance_norm) * evasion_strength * 15000
        
        return evasion_x, evasion_y
    
    def _execute_rush_strategy(self, black_boat: BlackBoat) -> Tuple[float, float]:
        """执行快速突破策略"""
        # 直接朝突破线冲刺
        rush_x = -20000  # 强力朝X=0方向移动
        rush_y = 0.0     # 保持Y方向
        
        return rush_x, rush_y


# 工厂函数
def create_black_agent(agent_type: str = "rule_based") -> BlackAgent:
    """
    创建黑方AI智能体
    
    Args:
        agent_type (str): AI类型 ("rule_based", "advanced")
        
    Returns:
        BlackAgent: 黑方AI实例
    """
    if agent_type in ["rule_based", "advanced"]:
        return RuleBasedBlackAgent()
    else:
        # 默认返回基于规则的AI
        return RuleBasedBlackAgent() 